import os
from PIL import ImageGrab
import cv2

class ImageTemplate():
    def find_image(self, template_image):
        # 构造图片保存的路径
        # base_path = os.path.join(os.getcwd(), '/source')
        base_path = os.path.join('D:\Workstation\PycharmProjects\SharedParking\source')
        # 截图文件的保存路径 = 图片的保存路径 + 截图名
        image_path = os.path.join(base_path, 'screen.png')
        image_path = os.path.abspath(image_path)
        # 截取当前屏幕图片并保存
        ImageGrab.grab().save(image_path, format='png')
        # 读取刚刚截的图，和准备好的模板图
        screen = cv2.imread(image_path)
        # 注意，模板图也要组合base_path路径
        template = cv2.imread(os.path.join(base_path, template_image))
        # 静态匹配，参数分别是截图，模板图，匹配算法，如果在截图中匹配到了模板图，则返回
        # 模板左顶点在每一个位置的匹配度，它是一个二维列表数据。
        result = cv2.matchTemplate(screen, template, cv2.TM_CCOEFF_NORMED)
        # 传入上一步得到的result，minMaxLoc方法的返回值是(最小值, 最大值, 最小值所在点坐标, 最大值所在点坐标)
        min, max, min_loc, max_loc = cv2.minMaxLoc(result)

        # 如果匹配成功，返回匹配到元素的中心坐标
        if max <= 0.9:
            return -1, -1
        else:
            x = max_loc[0] + int(template.shape[1] / 2)
            y = max_loc[1] + int(template.shape[0] / 2)

            return x, y

    # 断言
    def assertion(self, image):
        x, y = self.find_image(image)
        if x != -1 and y != -1:
            return True
        else:
            return False
